import os
import asyncio

from dotenv import load_dotenv
from neo4j._optional_deps import np

from lightrag import LightRAG, QueryParam
from lightrag.kg.shared_storage import initialize_pipeline_status
from lightrag.utils import setup_logger, EmbeddingFunc
from lightrag.llm.openai import openai_complete_if_cache, openai_embed

# 新增：加载 .env 文件中的环境变量
load_dotenv()

WORKING_DIR = "./dickens"  # 设置工作目录为 "dickens"

# 如果工作目录不存在，则创建该目录
if not os.path.exists(WORKING_DIR):
    os.mkdir(WORKING_DIR)


# 异步函数，用于调用 OpenAI 的嵌入模型获取文本的向量表示
async def embedding_func(texts: list[str]) -> np.ndarray:
    return await openai_embed(
        texts,  # 输入文本列表
        model=os.getenv("EMBEDDING_MODEL"),  # 从环境变量中获取嵌入模型的名称
        api_key=os.getenv("EMBEDDING_BINDING_API_KEY"),  # 从环境变量中获取嵌入模型的 API 密钥
        base_url=os.getenv("EMBEDDING_BINDING_HOST"),  # 从环境变量中获取嵌入模型的 API 基础 URL
    )


# 异步函数，用于获取嵌入向量的维度
async def get_embedding_dim():
    test_text = ["This is a test sentence."]  # 测试文本
    embedding = await embedding_func(test_text)  # 获取嵌入向量
    embedding_dim = embedding.shape[1]  # 获取嵌入向量的维度
    return embedding_dim

async def initialize_rag():
    embedding_dimension = await get_embedding_dim()  # 获取嵌入向量的维度
    print(f"Detected embedding dimension: {embedding_dimension}")
    print(os.getenv("LIGHTRAG_GRAPH_STORAGE"))
    # 创建 LightRAG 实例，传入工作目录、LLM 模型函数和嵌入模型函数
    rag = LightRAG(
        working_dir=WORKING_DIR,  # 设置工作目录
        # kv_storage=os.getenv("LIGHTRAG_KV_STORAGE"),
        # vector_storage=os.getenv("LIGHTRAG_VECTOR_STORAGE"),
        graph_storage=os.getenv("LIGHTRAG_GRAPH_STORAGE"),
        # doc_status_storage=os.getenv("LIGHTRAG_DOC_STATUS_STORAGE"),
        llm_model_func=llm_model_func,  # LLM 模型函数
        embedding_func=EmbeddingFunc(
            embedding_dim=embedding_dimension,  # 设置嵌入向量维度
            max_token_size=8192,  # 设置最大 token 数量
            func=embedding_func,  # 嵌入模型函数
        ),
    )

    # 初始化存储和管道状态
    await rag.initialize_storages()
    await initialize_pipeline_status()  # 初始化管道状态（可能是用于数据的处理和管理）

    return rag


# 异步函数，用于调用 LLM（大语言模型）进行文本生成
async def llm_model_func(
    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
) -> str:
    # 使用 openai_complete_if_cache 函数调用大语言模型（如果有缓存，则使用缓存）
    return await openai_complete_if_cache(
        os.getenv("LLM_MODEL"),  # 从环境变量中获取 LLM 模型的名称
        prompt,  # 输入提示（prompt）
        system_prompt=system_prompt,  # 系统提示（可选）
        history_messages=history_messages,  # 历史消息，用于上下文
        api_key=os.getenv("LLM_BINDING_API_KEY"),  # 从环境变量中获取 LLM API 密钥
        base_url=os.getenv("LLM_BINDING_HOST"),  # 从环境变量中获取 LLM API 基础 URL
        **kwargs,  # 其他参数
    )


async def main():
    # 初始化RAG实例
    rag = await initialize_rag()

if __name__ == "__main__":
    # 修改点：显式管理事件循环生命周期
    asyncio.run(main())
